from openai import OpenAI
import json  # 增加JSON解析，确保格式正确
from enum import Enum


class MyEnum(Enum):
    """Docstring for MyEnum."""
    enum = "enum"
    detectAndCls = "detectAndCls"
    QR = "QR"
    DM = "DM"
    CODE128 = "CODE128"
    
AreaType2index = {
    "enum": ["enum class AreaType", 0],
    "detectAndCls": ["class ModelFileDecryptionIV", 0],
    "QR": ["ModelFileDecryptionIV::QR_DECODE_MODEL_FILE_IV", 32],
    "DM": ["ModelFileDecryptionIV::DM_DECODE_MODEL_FILE_IV", 32],
    "CODE128": ["ModelFileDecryptionIV::CODE128_DECODE_MODEL_FILE_IV", 32],
}
client = OpenAI(
    api_key="sk-c4W8fCqxsjTEVkxRD7YIofvgfkFV81YQbuM09rbHW7zpol8A",
    base_url="https://api.moonshot.cn/v1",
)


def modify_code(user_prompt: str):
    completion = client.chat.completions.create(
        model="moonshot-v1-32k",  # 32k模型适合处理较长代码
        messages=[
            {
                "role": "system",
                "content": "你是擅长C++代码修改的助手，严格按照用户要求的位置和格式添加代码，返回完整内容，不省略任何部分。",
            },
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.1,  # 降低随机性，确保准确性
        response_format={"type": "json_object"},
        timeout=60,  # 延长超时时间
    )
    return completion.choices[0].message.content


def modify_AreaType(file_path: str, varname: str):
    # 1. 读取文件内容（增加异常处理，确保内容完整）
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件失败：{e}")
        return
    AreaType = "enum class AreaType"
    indexStart = content.index(AreaType)
    indexEnd = content.index("};", indexStart)
    contentAreaType = content[indexStart : indexEnd + 2]
    print(contentAreaType)
    user_prompt = f"这是全部的变量代码:{contentAreaType},这是需要添加的变量名称:{varname}，请添加变量并返回修改后的代码，键值为modify_code"
    result = modify_code(user_prompt)
    print(f"API返回结果：{result}")
    value = json.loads(result)["modify_code"]
    content = content.replace(contentAreaType, value)
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(content)
    # 添加成功AreaType
    print("添加成功AreaType")


def modify_QRCode(file_path: str, varname: str, iv: str, AreaType: str):
    # 1. 读取文件内容（增加异常处理，确保内容完整）
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件失败：{e}")
        return
    indexStart = content.index(AreaType2index[AreaType][0])
    indexEnd = content.index("};", indexStart)
    preStart = AreaType2index[AreaType][1]
    contentAreaType = content[indexStart - preStart : indexEnd + 2]
    print(contentAreaType)
    if AreaType == "detectAndCls":
        user_prompt = f"这是全部的变量代码:{contentAreaType},这是需要添加的变量名称:{varname},IV:{iv},请根据内容对该变量声明函数进行添加并只返回添加的代码,返回的代码键值key为modify_code,返回的变量名称键值为varname"
    else:
        user_prompt = f"这是全部的变量代码:{contentAreaType},这是需要添加的变量名称:{varname},IV:{iv},请添加变量并返回修改后的代码,返回的代码键值key为modify_code,返回的变量名称键值为varname"
    result = modify_code(user_prompt)
    print(f"API返回结果：{result}")
    value = json.loads(result)["modify_code"]
    varname = json.loads(result)["varname"]
    if AreaType == "detectAndCls":
        contentAreaTypeNew = contentAreaType[:-2]+"\n" + value + "\n" + contentAreaType[-2:]
        content = content.replace(contentAreaType, contentAreaTypeNew)
    else:
        content = content.replace(contentAreaType, value)
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(content)
    # 添加成功AreaType
    print(f"添加成功{AreaType}代码")
    return varname


def modify_Kotlin(
    file_path: str, varname: str, iv: str, AreaType: str, varnamePre: str = ""
):
    # 1. 读取文件内容（增加异常处理，确保内容完整）
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件失败：{e}")
        return
    user_prompt = f"这是全部的变量代码:{content},这是需要添加的变量名称:{varname},请根据内容对该变量声明函数进行添加并只返回添加的代码,返回的代码键值key为modify_code,同时将函数名称保存在function_name的键值里"
    result = modify_code(user_prompt)
    print(f"API返回结果:{result}")
    value = json.loads(result)["modify_code"]
    function_name = json.loads(result)["function_name"]
    # 寻找最后一个 "}" 的索引位置
    last_brace_index = content.rfind("}")
    # 在index位置插入value
    content = content[:last_brace_index] + value + "\n" + content[last_brace_index:]
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(content)
    try:
        with open(
            r"D:\data\android\demo\aidc\nativeLibs\app\src\main\cpp\cryption\model_file_decryption_iv.cpp",
            "r",
            encoding="utf-8",
        ) as f:
            content = f.read()
    except Exception as e:
        print(f"读取文件失败：{e}")
        return
    index_start = content.index('extern "C"')
    contentexternC = content[index_start:]
    user_prompt = f"这是一段JNI的代码,这是前面的示例代码:{contentexternC},这是需要添加的函数名称:{function_name},之前在ModelFileDecryptionIV中的变量名称为:{varnamePre},请根据上面的示例代码，返回函数名称的完整函数,返回的代码键值key为modify_code"
    result = modify_code(user_prompt)
    print(f"API返回结果：{result}")
    content = content + "\n" + json.loads(result)["modify_code"]
    with open(
        r"D:\data\android\demo\aidc\nativeLibs\app\src\main\cpp\cryption\model_file_decryption_iv.cpp",
        "w",
        encoding="utf-8",
    ) as f:
        f.write(content)
    # 添加成功AreaType
    print(f"添加成功{AreaType}代码")


def kimi(file_path, type: str, typeName: str):
    # 添加类型变量
    modify_AreaType(file_path, typeName)
    iv = "aef763d92205dd8b68ab520d71346e85"
    if type == "detectAndCls":
        varname = modify_QRCode(file_path, typeName, iv, AreaType=type)
        modify_Kotlin(
            r"D:\data\android\demo\aidc\nativeLibs\app\src\main\java\com\action\nativelibs\cryption/ModelFileDecryptionIV.kt",
            typeName,
            iv,
            AreaType=type,
            varnamePre=varname,
        )
    elif type == "QR":
        varname = modify_QRCode(file_path, typeName, iv, AreaType="QR")
        modify_Kotlin(
            r"D:\data\android\demo\aidc\nativeLibs\app\src\main\java\com\action\nativelibs\cryption/ModelFileDecryptionIV.kt",
            typeName,
            iv,
            AreaType="QR",
            varnamePre=varname,
        )
    elif type == "DM":
        varname = modify_QRCode(file_path, typeName, iv, AreaType="DM")
        modify_Kotlin(
            r"D:\data\android\demo\aidc\nativeLibs\app\src\main\java\com\action\nativelibs\cryption/ModelFileDecryptionIV.kt",
            typeName,
            iv,
            AreaType="DM",
            varnamePre=varname,
        )
    # elif type=="CODE128":
    #     modifyCode128(file_path, "QR_CODE_V13",iv)


# 调用函数（建议先备份原文件）
kimi(
    r"D:\data\android\demo\aidc\nativeLibs\app\src\main\cpp\cryption\model_file_decryption_iv.cpp",
    type=MyEnum.DM.value,
    typeName="DM_CODE_V14",
)